package com.example.tuike.common;

import com.example.tuike.entity.FetchTask;
import org.springframework.context.annotation.Bean;

import java.util.concurrent.PriorityBlockingQueue;
import java.util.concurrent.atomic.AtomicLong;

/*** @Author supersaiya
 * @Description TODO

 * @Date 2025/2/9 16:50

 * @Version 1.0*/
public class TokenBucketLimiter {
    private final long rate; // 每秒生成的令牌数
    private final long capacity; // 令牌桶的最大容量
    private final AtomicLong tokens; // 当前令牌数
    private final AtomicLong lastUpdate; // 上次更新令牌的时间（纳秒时间戳）

    public TokenBucketLimiter(long rate, long capacity) {
        this.rate = rate;
        this.capacity = capacity;
        this.tokens = new AtomicLong(capacity);
        this.lastUpdate = new AtomicLong(System.nanoTime());
    }

    public synchronized boolean acquire() {
        long now = System.nanoTime();
        long elapsed = now - lastUpdate.get();
        long newTokens = elapsed * rate / 1_000_000_000L; // 将纳秒转换为秒
        tokens.addAndGet(newTokens);
        tokens.set(Math.min(capacity, tokens.get()));
        lastUpdate.set(now);

        if (tokens.get() > 0) {
            tokens.decrementAndGet();
            return true;
        }
        return false;
    }
}